"""
A class to exercise the JT core functions.
"""
import JT
import unittest
import time
import itk
import vtk
import pylab
import numpy as np
from vtk.util import vtkImageExportToArray

from vtk.util import numpy_support


class Test(unittest.TestCase):
    def runTest(self):
        """
        Test out the ITK get metric procedure.
        """
#===============================================================================
#     Set up readers for the test
#===============================================================================
        #fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.png'
        fixedImageName = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'

        fixedImageReader = itk.ImageFileReader.IUS2.New()
        #vtkReader = vtk.vtkPNGReader()
        vtkReader = vtk.vtkMetaImageReader()

        fixedImageReader.SetFileName(fixedImageName)
        vtkReader.SetFileName(fixedImageName)
        
        vtkReader.Update()
        fixedImageReader.Update()
        
#===============================================================================
#    ITK -> python
#===============================================================================
        time1 = time.time()
        itk_pyimg = itk.PyBuffer.IUS2.GetArrayFromImage(fixedImageReader.GetOutput())
        time2 = time.time()
#===============================================================================
#    VTK -> python (shallow copy)
#===============================================================================
        vtk_pyimg = vtk.util.numpy_support.vtk_image_to_numpy(vtkReader.GetOutput())
        time3 = time.time()

#        pylab.ion()
#        pylab.figure()
#        pylab.subplot(211)
#        pylab.imshow(itk_pyimg)
#        pylab.subplot(212)
#        pylab.imshow(vtk_pyimg)
#        pylab.show()
#        
        t_itk_python = time2-time1
        t_vtk_python = time3-time2
        print "ITK->python: %E, VTK->python: %E"%(t_itk_python, t_vtk_python)
        return t_itk_python, t_vtk_python

if __name__ == "__main__":
    t=Test()
    t_itk = []
    t_vtk = []
    for i in range(100):
        t1, t2 = t.runTest()
        t_itk.append(t1)
        t_vtk.append(t2)
        
    print "---Averages---"
    t_itk = np.asarray(t_itk)
    t_vtk = np.asarray(t_vtk)
    print "ITK->python: %E, VTK->python: %E"%(t_itk.mean(), t_vtk.mean())